package com.alks.common.ratelimit.aspect;

import com.alks.common.ratelimit.annotation.RateLimit;
import com.alks.common.service.RedisService;
import com.alks.common.utils.IpAddressUtils;
import com.alks.entity.exception.RateLimitException;
import com.alks.entity.exception.ServiceErrorException;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.lang.reflect.Method;

@Slf4j
@Aspect
@Component
public class RateLimitAspect {

    @Autowired
    private RedisService redisService;

    @Pointcut("@annotation(com.alks.common.ratelimit.annotation.RateLimit)")
    public void rateLimitPointCut() {
    }

    @Around("rateLimitPointCut()")
    public Object around(ProceedingJoinPoint joinPoint) throws Throwable {
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();
        RateLimit rateLimit = method.getAnnotation(RateLimit.class);

        if (rateLimit != null) {
            String key = rateLimit.key();
            long limit = rateLimit.limit();
            long timeout = rateLimit.timeout();
            String message = rateLimit.message();

            if (key.isEmpty()) {
                // Default key based on IP address and method signature
                HttpServletRequest request = ((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getRequest();
                key = IpAddressUtils.getClientIpAddress(request) + ":" + method.getName();
            }

            String redisKey = "rate_limit:" + key;
            long currentCount = redisService.incr(redisKey, 1);

            if (currentCount == 1) {
                redisService.setOutTime(redisKey, timeout, rateLimit.timeUnit());
            }

            if (currentCount > limit) {
                log.warn("Rate limit exceeded for key: {}", redisKey);
                throw new RateLimitException(message);
            }
        }
        return joinPoint.proceed();
    }
}